from __future__ import division
from __future__ import print_function

import time
import argparse
from matplotlib import pyplot as plt
import numpy as np

import torch
import torch.nn.functional as F
import torch.optim as optim

from pygcn.utils import load_data, accuracy
from pygcn.models import GCN

# 命令行参数设置
parser = argparse.ArgumentParser()
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='Disables CUDA training.')
parser.add_argument('--fastmode', action='store_true', default=False,
                    help='Validate during training pass.')
parser.add_argument('--seed', type=int, default=42, help='Random seed.')
parser.add_argument('--epochs', type=int, default=500,
                    help='Number of epochs to train.')
parser.add_argument('--lr', type=float, default=0.02,
                    help='Initial learning rate.')
parser.add_argument('--weight_decay', type=float, default=5e-4,
                    help='Weight decay (L2 loss on parameters).')
parser.add_argument('--hidden', type=int, default=16,
                    help='Number of hidden units.')
parser.add_argument('--dropout', type=float, default=0.5,
                    help='Dropout rate (1 - keep probability).')

# 设置模型训练的超参数
args = parser.parse_args()
# 根据是否指定了 no_cuda和当前系统是否支持CUDA来确定是否使用GPU加速
args.cuda = not args.no_cuda and torch.cuda.is_available()

# 随机数种子设置--确保实验的可重复性
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

# 加载数据：邻接矩阵、特征矩阵、标签以及训练 验证 测试集的索引
adj, features, labels, idx_train, idx_val, idx_test = load_data()

# 调用GCN类构建GCN模型
model = GCN(nfeat=features.shape[1],
            nhid=args.hidden,
            nclass=labels.max().item() + 1,
            dropout=args.dropout)
# 使用Adam优化器进行参数更新（weight_decay表示正则化系数）
optimizer = optim.Adam(model.parameters(),
                       lr=args.lr, weight_decay=args.weight_decay)

# 如果过CUDA可用 将模型和数据移到GPU上
if args.cuda:
    model.cuda()
    features = features.cuda()
    adj = adj.cuda()
    labels = labels.cuda()
    idx_train = idx_train.cuda()
    idx_val = idx_val.cuda()
    idx_test = idx_test.cuda()


def train(epoch):
    t = time.time()
    model.train()# 将模型设置为训练模式。这时会启动训练中特定的操作例如dropout
    optimizer.zero_grad()# 固定操作，清楚优化器中保存的梯度
    output = model(features, adj)# 将输入特征和邻接矩阵传入模型得到输出
    # 负对数似然损失函数计算训练集上的损失。output[idx_train] 是模型在训练集上的预测结果，labels[idx_train] 是训练集上的真实标签。
    loss_train = F.nll_loss(output[idx_train], labels[idx_train])
    # 计算训练集上的准确率，accuracy 是一个自定义函数
    acc_train = accuracy(output[idx_train], labels[idx_train])
    loss_train.backward()# 反向传播计算梯度，根据损失函数和模型的参数计算各参数的梯度
    optimizer.step()# 根据优化器设置更新模型参数

    if not args.fastmode:
        # 若没有设置快速模式将模型设置为评估模式。会关闭dropout，确保评估的一致性
        model.eval()
        output = model(features, adj)

    loss_val = F.nll_loss(output[idx_val], labels[idx_val])# 计算验证集上的损失
    acc_val = accuracy(output[idx_val], labels[idx_val])# 计算准确率
    print('Epoch: {:04d}'.format(epoch+1),
          'loss_train: {:.4f}'.format(loss_train.item()),
          'acc_train: {:.4f}'.format(acc_train.item()),
          'loss_val: {:.4f}'.format(loss_val.item()),
          'acc_val: {:.4f}'.format(acc_val.item()),
          'time: {:.4f}s'.format(time.time() - t))
    return loss_val.item()

def test():
    model.eval()
    output = model(features, adj)
    loss_test = F.nll_loss(output[idx_test], labels[idx_test])
    acc_test = accuracy(output[idx_test], labels[idx_test])
    print("Test set results:",
          "loss= {:.4f}".format(loss_test.item()),
          "accuracy= {:.4f}".format(acc_test.item()))


# 模型训练循环
losses = []
t_total = time.time()
for epoch in range(args.epochs):
    loss = train(epoch)
    losses.append(loss)
print("Optimization Finished!")
print("Total time elapsed: {:.4f}s".format(time.time() - t_total))
# 模型测试
test()
# 绘制loss曲线
plt.figure()
plt.plot(losses)
plt.show()
